import os
# os.environ['CUDA_VISIBLE_DEVICES'] = "0"
from keras.callbacks import LearningRateScheduler, ModelCheckpoint, TensorBoard
from keras.optimizers import SGD
from keras.preprocessing.image import ImageDataGenerator
from Models import Multi_models


#  训练过程中lr改变规律
def lr_decay(epoch):
    if epoch <= 20:
        return 0.1
    if epoch <= 40:
        return 0.01
    else:
        return 0.001


class train():

    def __init__(self):
        self.train_path = './data/train_data/'
        self.val_path = './data/val_data/'
        self.test_image_path = './data/test1/test1/'
        self.BATCH_SIZE = 32
        self.WEIGHTS_PATH = 'logs/best_mobile.h5'
        self.max_Epochs = 60

    def model_train(self, model, is_loadweights, which_model):
        history = None
        # 对训练数据进行数据增强
        train_datagen = ImageDataGenerator(
                                            # width_shift_range=0.2,
                                            # height_shift_range=0.2,
                                            rescale=1./255,
                                            # shear_range=0.2,
                                            zoom_range=0.2,
                                            horizontal_flip=True,
                                            fill_mode='nearest')
        # train_datagen = ImageDataGenerator(rescale=1./255,
        #                                    fill_mode='nearest',
        #                                    rotation_range=10)

        # 验证集不做增强
        val_datagen = ImageDataGenerator(rescale=1. / 255)
        # 从本地取出训练数据集
        train_generator = train_datagen.flow_from_directory(self.train_path,
                                                            target_size=(128, 128),
                                                            batch_size=self.BATCH_SIZE, color_mode='grayscale',
                                                            class_mode='categorical')
        val_generator = val_datagen.flow_from_directory(self.val_path,
                                                        target_size=(128, 128),
                                                        batch_size=self.BATCH_SIZE, color_mode='grayscale',
                                                        class_mode='categorical')
        lr = LearningRateScheduler(lr_decay)
        mdcheck = ModelCheckpoint(self.WEIGHTS_PATH, monitor='val_acc', save_best_only=True)
        td = TensorBoard(log_dir='./data/tensorboard_log/')
        # 是否导入训练过的权重参数
        if is_loadweights:
            if os.path.isfile(self.WEIGHTS_PATH):
                model.load_weights(self.WEIGHTS_PATH, by_name=True)
                print('model have load pre weights of hanzi image !!')
            else:
                print('model not load weights!!')
        else:
            print('not load weights model')

        sgd = SGD(lr=0.1, momentum=0.9, decay=5e-4, nesterov=True)
        if which_model == "cnn":
            print("model compile!!")
            model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=['accuracy'])
            print("model training!!")
            history = model.fit_generator(train_generator, steps_per_epoch=34000 // self.BATCH_SIZE,
                                          epochs=self.max_Epochs, validation_data=val_generator,
                                          validation_steps=6000 // self.BATCH_SIZE, callbacks=[lr, mdcheck, td])
            model.save_weights("logs/last_cnn.h5")
        elif which_model == "mobile":
            print("model compile!!")
            model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=['accuracy'])
            print("model training!!")
            history = model.fit_generator(train_generator, steps_per_epoch=34000 // self.BATCH_SIZE,
                                          epochs=self.max_Epochs, validation_data=val_generator,
                                          validation_steps=6000 // self.BATCH_SIZE, callbacks=[lr, mdcheck, td])
        elif which_model == "densenet":
            print("model compile!!")
            model.compile(optimizer=sgd, loss='categorical_crossentropy', metrics=['accuracy'])
            print("model training!!")
            history = model.fit_generator(train_generator, steps_per_epoch=34000 // self.BATCH_SIZE,
                                          epochs=self.max_Epochs, validation_data=val_generator,
                                          validation_steps=6000 // self.BATCH_SIZE, callbacks=[lr, mdcheck, td])

        return history


if __name__ == '__main__':

    # mmodels = Multi_models(CLASS=100, size=128)
    # simple_model = mmodels.Dense_mdoel()
    # print(simple_model.summary())
    #
    # print("=====start train image of epoch=====")
    # train = train()
    # model_history = train.model_train(simple_model, True, which_model="densenet")

    # mmodels = Multi_models(CLASS=100, size=128)
    # simple_model = mmodels.cnn_model()
    # print(simple_model.summary())
    #
    # print("=====start train image of epoch=====")
    # train = train()
    # model_history = train.model_train(simple_model, True, which_model='cnn')

    mmodels = Multi_models(CLASS=100, size=128)
    simple_model = mmodels.mobile()
    print(simple_model.summary())

    print("=====start train image of epoch=====")
    train = train()
    model_history = train.model_train(simple_model, False, which_model='mobile')
